import { RdTypeEnum, WorkFlowTypeEnum, type TemplateType } from "@/utils/enums";
import BaseAInstance, { type IBaseAIParams } from "./useBaseAInstance";
import { type TemplateChooseObjStore } from "@/store/indexType";
import { type ImageParametersType } from "@/store/modules/workflowStore";
import { getUUID } from "@/utils/common";

import {
  oldBuildingReBuildPrompt3,
  //   pictureCreateDraw,
  pictureCreateByLinePrompt3,
  modelRenderPrompt3,
  inspirationMovePrompt3,
  roomReBuildPrompt3,
  cfg_scale_1,
  sampling_steps_1,
  paramsDic,
  templateSelectDic,
  p4,
  np4,
  type PresetItem,
} from "@/utils/dicts";
import { imageToImage, textToImage } from "@/api/aiDraw";

//  设置一些奇奇怪怪的默认参数
const WORKFLOWTYPES = [
  {
    label: "线稿生成",
    value: WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE,
    // 不知名参数
    baseTag: "",
    styleTag: "",
    noOpenBaseTag: "workflow_lineart",
    noOpenStyleTag: "workflow_lineart_IPA",
  },
  {
    label: "模型渲染",
    value: WorkFlowTypeEnum.MODEL_RENDER,
    baseTag: "workflow_model_render_img2img",
    styleTag: "workflow_model_render_img2img_IPA",
    noOpenBaseTag: "workflow_model_render",
    noOpenStyleTag: "workflow_model_render_IPA",
  },
  {
    label: "旧房改造",
    value: WorkFlowTypeEnum.OLD_BUILDING_REBUILD,
    baseTag: "workflow_facade_renovation_img2img",
    styleTag: "workflow_facade_renovation_img2img_IPA",
    noOpenBaseTag: "workflow_facade_renovation_txt2img",
    noOpenStyleTag: "workflow_facade_renovation_txt2img_IPA",
  },
  {
    label: "房间装修",
    value: WorkFlowTypeEnum.ROOM_REBUILD,
    // 不知名参数
    baseTag: "workflow_room_img2img_furniture", // （带家具+开原图+不上传风格参考图）
    styleTag: "workflow_room_img2img_furniture_IPA", // （带家具+开原图+上传风格参考图）
    noOpenBaseTag: "workflow_room_txt2img_furniture", // （带家具+关原图+不上传风格参考图）
    noOpenStyleTag: "workflow_room_txt2img_furniture_IPA", // （带家具+关原图+上传风格参考图）
    roughTag: "workflow_room_txt2img", // （毛坯房+不上传风格参考图）
    roughStyleTag: "workflow_room_txt2img_IPA", // （毛坯房+上传风格参考图）
  },
  {
    label: "灵感迁移",
    value: WorkFlowTypeEnum.INSPIRATION_MOVE,
    // 不知名参数
    baseTag: "workflow_transfer",
    styleTag: "",
    noOpenBaseTag: "",
    noOpenStyleTag: "",
  },
];

export const drawObj = {
  [WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE]: pictureCreateByLinePrompt3, // 线稿生成
  [WorkFlowTypeEnum.MODEL_RENDER]: modelRenderPrompt3, // 模型渲染
  [WorkFlowTypeEnum.INSPIRATION_MOVE]: inspirationMovePrompt3, // 灵感迁移
  [WorkFlowTypeEnum.ROOM_REBUILD]: roomReBuildPrompt3, // 房间装修
  [WorkFlowTypeEnum.OLD_BUILDING_REBUILD]: oldBuildingReBuildPrompt3, // 旧房改造
};
class WorkflowGenerator extends BaseAInstance {
  private workflowParams: Omit<IWorkflowParams, keyof IBaseAIParams>; // api需要的
  private frontParams: IFrontParams; // 页面中需要的参数

  constructor(params: IFrontParams, workflowParam: Partial<IWorkflowParams> = {}) {
    super(workflowParam);
    this.frontParams = {
      ...params,
    };
    this.workflowParams = this.generateDefaultParams();
  }

  // 生成默认参数
  private generateDefaultParams = () => {
    // 不依赖 frontParams
    const reqWorkflowParams = {
      atmosphere: "通用",
      visualAngle: "通用",
      environment: "通用",
      batch_size: 1,
      force_task_id: getUUID(),
      seed: Math.floor(Math.random() * (262144 - 65536)) + 65536,
    };

    return reqWorkflowParams as Omit<IWorkflowParams, keyof IBaseAIParams>;
  };
  // 生成图片
  generate = async () => {
    await this.handleParams();
    if (this.frontParams.workflowType === WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE) {
      await textToImage({ ...this.workflowParams, ...this.params });
    } else if (this.frontParams.workflowType === WorkFlowTypeEnum.INSPIRATION_MOVE) {
      await imageToImage({ ...this.workflowParams, ...this.params });
    } else {
      if (
        (this.frontParams.workflowType !== WorkFlowTypeEnum.ROOM_REBUILD ||
          (this.frontParams.workflowType === WorkFlowTypeEnum.ROOM_REBUILD &&
            this.frontParams.imageParameters.rdType === RdTypeEnum.DETAILED)) &&
        this.frontParams.imageParameters.originStrength.checked
      ) {
        await imageToImage({ ...this.workflowParams, ...this.params });
      } else {
        await textToImage({ ...this.workflowParams, ...this.params });
      }
    }
  };

  // 处理参数
  public handleParams = async () => {
    const frontParams = this.frontParams;
    if (frontParams) {
      const { batchSize, cueWordObj, imageObj, imageParameters, templateChooseObj, workflowType } =
        frontParams;
      let denoisingStrength = undefined;
      if (workflowType === WorkFlowTypeEnum.INSPIRATION_MOVE) {
        denoisingStrength = 1 - imageParameters.similarity;
      } else if (workflowType === WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE) {
        denoisingStrength = undefined;
      } else {
        denoisingStrength = imageParameters.originStrength.checked
          ? 1 - imageParameters.originStrength.value
          : undefined;
      }

      let way = "工作流-" + WORKFLOWTYPES.find((item) => item.value === workflowType)?.label;

      // 区分工作流类型
      const linePrompts = drawObj[workflowType].filter(
        (item) => item.templateNum === templateChooseObj.currentTemplate.value,
      );

      const linePrompt =
        linePrompts.length > 1
          ? linePrompts.find((item) => item.type === imageParameters.rdType)
          : linePrompts[0];

      // 这是个奇怪的字段，单独处理
      const tag = this.handleTag(workflowType, imageParameters, imageObj);

      /* 奇怪的正/反 提示词 */
      const miniIndustryObj = this.getcurrenTTemplateObjByKey(
        templateChooseObj.currentTemplate.value as number,
        templateChooseObj.businessOption.value as string,
        templateChooseObj.subdivision.value as string,
      );
      const p2 = this.getP2Prompt(
        templateChooseObj.currentTemplate.value as number,
        this.workflowParams.visualAngle,
        this.workflowParams.environment,
        this.workflowParams.atmosphere,
      );
      const prompt =
        miniIndustryObj?.p1 +
        (`${p2}` ? `,${p2}` : ``) +
        (cueWordObj.prompt ? `,${cueWordObj.prompt},` : `,`) +
        miniIndustryObj?.p3 +
        `,${p4}`;
      const negative_prompt =
        (cueWordObj.negativePrompt ? `${cueWordObj.negativePrompt},` : ``) + `${np4}`;
      /* ---- 奇怪的正/反 提示词 --- */
      // 图片宽高处理
      const info = await this.getImgSize(imageObj.baseImgUrl);
      const calcWH = this.getCalcOriginWh(
        {
          width: info.width,
          height: info.height,
        },
        true,
      );

      this.workflowParams = {
        ...this.workflowParams,
        batch_size: batchSize || 1,
        chooseTemplate: templateChooseObj.currentTemplate.value as number,
        industryKey: templateChooseObj.businessOption.value as string,
        industryMinKey: templateChooseObj.subdivision.value as string,
        ui_prompt: cueWordObj.prompt || "",
        ui_negative_prompt: cueWordObj.negativePrompt || "",
        prompt_optimation: cueWordObj.promptOptimation,
        rdType: imageParameters.rdType,
        denoising_strength: denoisingStrength,
        way: way,
        cfg_scale: linePrompt?.cfg_scale || cfg_scale_1,
        steps: linePrompt?.samplingSteps || sampling_steps_1,
        tag,
        prompt,
        negative_prompt,
        width: calcWH.width || 1024,
        height: calcWH.height || 1024,
        override_settings: {
          sd_model_checkpoint:
            paramsDic[templateChooseObj.subdivision.value as keyof typeof paramsDic]?.modelParams ||
            "",
        },
        alwayson_scripts: {
          controlnet: {
            args: [],
          },
        },
        init_images_uuid: "",
      };

      //   处理图片
      await this.handleImages(
        workflowType,
        imageObj,
        imageParameters,
        linePrompt as Partial<PresetItem>,
      );
    }
  };

  // 处理图片信息
  private handleImages = async (
    workflowType: WorkFlowTypeEnum,
    imageObj: { baseImgUrl: string; styleImgUrl: string },
    imageParameters: ImageParametersType,
    linePrompt: Partial<PresetItem>,
  ) => {
    return new Promise<void>(async (resolve, reject) => {
      try {
        let args: Args[] = [];
        switch (workflowType) {
          case WorkFlowTypeEnum.INSPIRATION_MOVE:
            const data = await this.uploadImages(await this.imageToBase64(imageObj.baseImgUrl));
            this.workflowParams.init_images_uuid = data.imageUuid;
            break;
          case WorkFlowTypeEnum.MODEL_RENDER:
          case WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE:
          case WorkFlowTypeEnum.OLD_BUILDING_REBUILD:
          case WorkFlowTypeEnum.ROOM_REBUILD:
            if (imageObj.baseImgUrl) {
              args.push(
                this.CNArgbuilder(
                  "canny",
                  "diffusers_xl_canny_full [2b69fca4]",
                  linePrompt.cannyWeight || 1,
                  await this.imageToBase64(imageObj.baseImgUrl),
                ),
              );
              if (
                workflowType !== WorkFlowTypeEnum.ROOM_REBUILD ||
                (workflowType === WorkFlowTypeEnum.ROOM_REBUILD &&
                  imageParameters.rdType === RdTypeEnum.DETAILED)
              ) {
                args.push(
                  this.CNArgbuilder(
                    "depth_midas",
                    "diffusers_xl_depth_full [2f51180b]",
                    linePrompt.depthWeight,
                    await this.imageToBase64(imageObj.baseImgUrl),
                  ),
                );
              }

              if (workflowType !== WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE) {
                const data = await this.uploadImages(await this.imageToBase64(imageObj.baseImgUrl));
                this.workflowParams.init_images_uuid = data.imageUuid;
              }
            }
            if (imageObj.styleImgUrl) {
              args.push(
                this.CNArgbuilder(
                  "ip-adapter_clip_g",
                  "ip-adapter_xl [4209e9f7]",
                  imageParameters.similarity,
                  await this.imageToBase64(imageObj.styleImgUrl),
                ),
              );
            }

            //   上传图片，拿到uuid
            const CNArgsRes = await Promise.all(
              args.map((item) => {
                return this.uploadImages(item.input_image as Base64URLString);
              }),
            );
            args = args.map((item, index) => {
              item.UUID = CNArgsRes[index].imageUuid;
              item.input_image = undefined;
              return item;
            });

            this.workflowParams.alwayson_scripts.controlnet.args = args;
            break;
          default:
            break;
        }
        resolve();
      } catch (error) {
        reject(error);
      }
    });
  };

  // 处理奇怪参数 - tag
  private handleTag = (
    workflowType: WorkFlowTypeEnum,
    imageParameters: ImageParametersType,
    imageObj: {
      baseImgUrl: string;
      styleImgUrl: string;
    },
  ) => {
    let tag = "";
    const currentTagObj = WORKFLOWTYPES.find((item) => item.value === workflowType);
    switch (workflowType) {
      case WorkFlowTypeEnum.INSPIRATION_MOVE:
        tag = currentTagObj?.baseTag || "";
        break;
      case WorkFlowTypeEnum.PICTURE_CREATE_BY_LINE:
        tag =
          (imageObj.styleImgUrl ? currentTagObj?.noOpenStyleTag : currentTagObj?.noOpenBaseTag) ||
          "";
      case WorkFlowTypeEnum.MODEL_RENDER:
      case WorkFlowTypeEnum.OLD_BUILDING_REBUILD:
        if (imageObj.styleImgUrl) {
          tag =
            (imageParameters.originStrength.checked
              ? currentTagObj?.styleTag
              : currentTagObj?.noOpenStyleTag) || "";
        } else {
          tag =
            (imageParameters.originStrength.checked
              ? currentTagObj?.baseTag
              : currentTagObj?.noOpenBaseTag) || "";
        }
      case WorkFlowTypeEnum.ROOM_REBUILD:
        if (imageObj.styleImgUrl) {
          if (imageParameters.rdType === RdTypeEnum.DETAILED) {
            tag =
              (imageParameters.originStrength.checked
                ? currentTagObj?.styleTag
                : currentTagObj?.noOpenStyleTag) || "";
          } else {
            tag = currentTagObj?.roughStyleTag || "";
          }
        } else {
          if (imageParameters.rdType === RdTypeEnum.DETAILED) {
            tag =
              (imageParameters.originStrength.checked
                ? currentTagObj?.baseTag
                : currentTagObj?.noOpenBaseTag) || "";
          } else {
            tag = currentTagObj?.roughTag || "";
          }
        }
    }

    return tag;
  };

  // 从PC端继承过来的方法
  private getcurrenTTemplateObjByKey = (
    chooseTemplate: number,
    industryKey: string,
    industryMinKey: string,
  ) => {
    if (industryKey && industryMinKey && chooseTemplate) {
      const chooseTemplates = templateSelectDic[chooseTemplate as keyof typeof templateSelectDic];
      const industrys = chooseTemplates.type.find((item) => item.key === industryKey);
      const miniIndustrys = industrys?.list.find((item) => item.key === industryMinKey);
      return miniIndustrys;
    }
    return null;
  };

  // 从PC端继承过来的方法
  private getP2Prompt = (
    currentChooseTemplate: number,
    visualAngle: string,
    environment: string,
    atmosphere: string,
  ) => {
    const p2Array =
      templateSelectDic[currentChooseTemplate as keyof typeof templateSelectDic].p2Array || [];
    const visualAngleList = p2Array.find((item) => item.label == "视角").options.flat();
    const environmentList = p2Array.find((item) => item.label == "环境").options.flat();
    const atmosphereList = p2Array.find((item) => item.label == "氛围").options.flat();
    // 通用没有p2所以没找到没关系
    const p2_1 = visualAngleList.find((item: { name: string }) => item.name == visualAngle)?.p2;
    const p2_2 = environmentList.find((item: { name: string }) => item.name == environment)?.p2;
    const p2_3 = atmosphereList.find((item: { name: string }) => item.name == atmosphere)?.p2;
    const p2 = [p2_1, p2_2, p2_3].filter((item) => !!item);
    return `${p2.join(",")}`;
  };
}

export default WorkflowGenerator;

export interface IWorkflowParams extends IBaseAIParams {
  atmosphere: string;
  visualAngle: string;
  environment: string;
  batch_size: number; // 生成数量
  baseImageUUID?: string;
  init_images_uuid?: string;
  cfg_scale: number;
  chooseTemplate: TemplateType;
  force_task_id: string;
  width: number; // 原图宽
  height: number; // 原图高
  industryKey: string; // 模板-业态
  industryMinKey: string; // 模板-细分类型
  prompt: string; // 正向提示词
  negative_prompt: string; // 反向提示词
  prompt_optimation: boolean; // 提示词优化
  rdType: RdTypeEnum; // 模型粒度
  denoising_strength?: number | null;
  seed: number; // 随机种子
  steps: number; // 采样步数
  tag: string;
  ui_prompt: string; // 正向提示词
  ui_negative_prompt: string; // 反向提示词
  way: string;
  override_settings: {
    sd_model_checkpoint: string;
  };
  alwayson_scripts: {
    controlnet: {
      args: Args[] | [];
    };
  };
}

export interface Args {
  model: string;
  module: string;
  processor_res: number;
  trigg_symb: boolean;
  input_image: string | undefined;
  weight: number;
  guidance_start: number;
  guidance_end: number;
  control_mode: string;
  UUID: string | null; // 上传后得到UUID
}

export interface IFrontParams {
  workflowType: WorkFlowTypeEnum; // 工作流类型
  templateChooseObj: TemplateChooseObjStore; // 模板选择对象
  imageObj: {
    baseImgUrl: string;
    styleImgUrl: string;
  }; // 图片对象
  imageParameters: ImageParametersType; // 图片参数
  cueWordObj: {
    prompt: string;
    negativePrompt: string;
    promptOptimation: boolean;
  }; // 提示词
  batchSize: number; // 生成数量
}
